# 读取json文件
import json
import numpy as np
from matplotlib import pyplot as plt

from algo import visualize_results


def visualize_results_copy(results, PLOT_IMAGE):
    """可视化延迟对比（横坐标为总流量，重复值取首个）"""
    plt.figure(figsize=(14, 7))

    # 数据准备与去重处理
    scale_keys = sorted(results.keys(), key=lambda x: float(x.split('_')[1]))

    seen_traffic = set()  # 用于记录已出现的流量值
    filtered_indices = []  # 保留的数据点索引
    base_traffic_volumes = []

    # 第一次遍历：去重处理
    for idx, key in enumerate(scale_keys):
        traffic = np.sum(results[key]["base_traffic"])  # 使用基础流量作为横坐标
        if traffic not in seen_traffic:
            seen_traffic.add(traffic)
            base_traffic_volumes.append(traffic)
            filtered_indices.append(idx)

    # 算法样式配置
    algorithms = {
        "full_update": {"color": "#1f77b4", "label": "全更新算法"},
        "semi_update": {"color": "#ff7f0e", "label": "半更新算法"},
        "hybrid_static": {"color": "#d62728", "label": "静态混合"}
    }

    # 绘制曲线
    for algo, style in algorithms.items():
        # 从完整数据中筛选去重后的数据点
        max_delays = [results[scale_keys[i]]["total_delays"][algo] for i in filtered_indices]

        plt.plot(
            base_traffic_volumes,
            max_delays,
            linewidth=2,
            marker='o',  # 统一添加圆形标记
            markersize=8,
            color=style["color"],
            label=style["label"]
        )

    # 图表美化
    plt.title("不同算法最大延迟随基础流量变化趋势",
              fontsize=14, fontfamily='SimHei')
    plt.xlabel("allreduce基础流量总量 (GB)", fontsize=12, fontfamily='SimHei')
    plt.ylabel("总共延迟 (秒)", fontsize=12, fontfamily='SimHei')

    # 智能坐标轴格式化
    max_traffic = max(base_traffic_volumes) if base_traffic_volumes else 0
    if max_traffic > 10000:
        plt.gca().xaxis.set_major_formatter(
            plt.FuncFormatter(lambda x, _: f"{x / 1000:.0f}k"))
        plt.xlabel("基础流量总量 (×1000 GB)")
    else:
        plt.gca().xaxis.set_major_formatter(
            plt.FuncFormatter(lambda x, _: f"{int(x)}"))

    plt.grid(True, alpha=0.3)
    plt.legend(fontsize=10, loc='upper left')

    # 添加数据标签（仅显示前3个点示例）
    for algo, style in algorithms.items():
        max_delays = [results[scale_keys[i]]["total_delays"][algo] for i in filtered_indices]
        for x, y in list(zip(base_traffic_volumes, max_delays))[:3]:  # 只标注前3个点
            plt.text(x, y + 0.1, f'{y:.1f}s',
                     fontsize=8, color=style["color"],
                     ha='center', va='bottom')

    plt.tight_layout()
    plt.savefig(PLOT_IMAGE, dpi=300, bbox_inches='tight')
    plt.show()
# 保存结果
with open("algorithm_results_allreduce_16_2.json", 'r') as f:
    analysis_results = json.load(f)
PLOT_IMAGE="total_allreduce_delay_comparison.png"
# 可视化
visualize_results_copy(analysis_results, PLOT_IMAGE)